import { Tensor, TensorShape } from "./Tensor";

export enum LayerType {
    Nop = 0,
    Dense = 1,
    MatMul = 2,
    Dense3 = 3,
    Conv2D = 20,
    DepthwiseConv2D = 21,
    Conv2DTrans = 22,
    Upsample2D = 23,
    MaxPool2D = 25,
    AvgPool2D = 26,
    GlobalMaxPool2D = 27,
    GlobalAvgPool2D = 28,
    Border2D = 29,
    Conv3D = 30,
    Conv3DTrans = 32,
    Upsample3D = 33,
    MaxPool3D = 35,
    AvgPool3D = 36,
    GlobalMaxPool3D = 37,
    GlobalAvgPool3D = 38,
    Border3D = 39,
    Activation = 50,
    ScaleBias = 51,
    Normalization = 52,
    LRN = 53,
    Dropout = 60,
    RandomNormal = 64,
    RandomUniform = 65,
    Multinomial = 66,
    OneHot = 67,
    TopKIndices = 68,
    TopKValues = 69,
    NonZero = 70,
    Range = 71,
    RoiAlign = 72,
    Add = 100,
    Sub = 101,
    Mul = 102,
    Div = 103,
    Pow = 104,
    Min = 110,
    Max = 111,
    Mean = 112,
    ReduceL1 = 120,
    ReduceL2 = 121,
    ReduceLogSum = 122,
    ReduceLogSumExp = 123,
    ReduceMax = 124,
    ReduceMean = 125,
    ReduceMin = 126,
    ReduceProd = 127,
    ReduceSum = 128,
    ReduceSumSquare = 129,
    Greater = 140,
    GreaterEqual = 141,
    Less = 142,
    LessEqual = 143,
    Equal = 144,
    LogicalOr = 145,
    LogicalAnd = 146,
    LogicalNot = 147,
    LogicalXor = 148,
    Where = 149,
    Sign = 150,
    Pad = 159,
    Pad2DReflect = 160,
    Pad2DSymmetric = 161,
    Pad2DEdge = 162,
    ArgMax = 163,
    ArgMin = 164,
    ConstantOfShape = 199,
    Flatten = 200,
    Reshape = 201,
    Transpose = 202,
    Squeeze = 203,
    Unsqueeze = 204,
    Gather = 205,
    DepthToSpace = 206,
    SpaceToDepth = 207,
    Expand = 208,
    Resample2D = 209,
    Concat = 210,
    StridedSlice = 211,
    Tile = 212,
    Shape = 213,
    NonMaxSuppression = 214,
    LSTM = 215,
    ScatterND = 216,
    Load = 255
}

export enum ActivationType {
    None = 0,
    Relu = 1,
    Softmax = 2,
    Tanh = 3,
    Sigmoid = 4,
    Elu = 5,
    Relu6 = 6,
    LeakyRelu = 7,
    Selu = 8,
    Swish = 9,
    LogSoftmax = 10,
    Softplus = 11,
    Softsign = 12,
    PRelu = 13,
    Hardmax = 20,
    HardSigmoid = 21,
    Abs = 100,
    Neg = 101,
    Ceil = 102,
    Clip = 103,
    Floor = 104,
    Round = 105,
    Reciprocal = 110,
    Sqrt = 111,
    Pow = 112,
    Exp = 113,
    Log = 114,
    Acos = 200,
    Acosh = 201,
    Asin = 202,
    Asinh = 203,
    Atan = 204,
    Atanh = 205,
    Cos = 206,
    Cosh = 207,
    Sin = 208,
    Sinh = 209,
    Tan = 210,
    Erf = 211
}

export enum AutoPad {
    NotSet = 1,
    Valid = 0,
    SameUpper = -1,
    SameLower = -2
}

export enum PadMode {
    Constant = 0,
    Reflect = 1,
    Edge = 2,
    Symetric = 3
}

export enum DepthToSpaceMode {
    DCR,
    CRD
}

export enum ScatterNDReductionMode {
    None = 0,
    Add = 1,
    Mul = 2
}

export interface DataSet {
    name: string;
    shape: TensorShape;
    offset: number;
    itemSizeInBytes: number;
    length: number;
}

export class Layer {
    public name: string;
    public type: LayerType;
    public flags: number = 0;
    public activation: ActivationType;
    public pad: number[];
    public stride: number[];
    public pool: number[];
    public axis: number;
    public alpha: number;
    public beta: number;
    public inputs: string[];
    public outputs: string[];
    public axes: number[];
    public datasets: DataSet[];
    public weights: number[];

    constructor(layerName: string, layerType: LayerType, activationType: ActivationType = ActivationType.None) {
        this.name = layerName;
        this.type = layerType;
        this.activation = activationType;
        this.pad = [];
        this.stride = [];
        this.pool = [];
        this.axis = -1;
        this.alpha = 1.0;
        this.beta = 0.0;
        this.inputs = [];
        this.outputs = [];
        this.axes = [];
        this.datasets = [];
        this.weights = [];
    }

    public toString(): string {
        return `name:${this.name}, activation:${this.activation}, inputs:[${this.inputs.join(",")}], 
        pad:[${this.pad.join(",")}], stride:[${this.stride.join(",")}], pool:[${this.pool.join(",")}],
         alpha:${this.alpha}, beta:${this.beta}, axis:${this.axis}, 
         weights:[${this.datasets.map(ds => `${ds.name} ${ds.shape}`).join(",")}]`
            .replace(`${this.name}/`, "").replace(`${this.name} `, "")
            .replace("activation:None, ", "").replace("inputs:[], ", "")
            .replace("pad:[], ", "").replace("stride:[], ", "")
            .replace("stride:[1,1], ", "").replace("pool:[], ", "")
            .replace("alpha:1, ", "").replace("beta:0, ", "")
            .replace("axis:-1, ", "").replace("weights:[]", "");
    }

    public dataSetToTensor(index: number): Tensor | null {
        if (index < this.datasets.length) {
            const ds = this.datasets[index];
            const tensor = new Tensor();
            tensor.shape = ds.shape;
            tensor.name = ds.name;
            tensor.cache = this.weights;
            return tensor
        }
        return null;
    }

    public applyTensorToDataSet(X: Tensor, index: number): void {
        if (index < this.datasets.length) {
            const ds = this.datasets[index];
            ds.shape = X.shape;
            this.weights = Array.from(X.cache);
            this.datasets[index] = ds;
        }
    }
}


export interface Input {
    name: string;
    shape: number[];
    rank?: number;
}

export interface Memory {
    shape: TensorShape;
    input: string;
    output: string;
}

export enum DataType {
    Float32 = 0,
    Float16 = 1,
    Int32 = 2,
    UInt8 = 3,
    Int8 = 4,
    String = 5
}

export class ImporterWarning {
    public readonly message: string;
    public readonly layerName: string;

    constructor(layer: string, msg: string) {
        this.message = msg;
        this.layerName = layer;
    }
}


export class Model {
    public static readonly Version: number = 20;
    public static readonly LastVersionWithout8DSupport: number = 16;
    public static readonly LastVersionWithoutWeightsAlignmentSupport: number = 18;
    public static readonly WeightsAlignment: number = 16;

    public layout: string = '';
    public inputs: Input[] = [];
    public outputs: string[] = [];
    public memories: Memory[] = [];
    public layers: Layer[] = [];

    public IrSource: string = 'Script';
    public IrVersion: string = 'NA';
    public ProducerName: string = 'Script';
    public Warnings: ImporterWarning[] = [];
    public Metadata: { [key: string]: string } = {};

    public ShallowCopy(): Model {
        const model = new Model();
        model.inputs = this.inputs.slice();
        model.outputs = this.outputs.slice();
        model.memories = this.memories.slice();
        model.layers = this.layers.slice();

        model.IrSource = this.IrSource;
        model.IrVersion = this.IrVersion;
        model.ProducerName = this.ProducerName;
        model.Warnings = this.Warnings.slice();
        model.Metadata = { ...this.Metadata };
        return model;
    }

    public toString(): string {
        let totalUniqueWeights = 0;
        for (const layer of this.layers) {
            for (const dataset of layer.datasets) {
                totalUniqueWeights += dataset.length;
            }
        }

        return `inputs: [${this.inputs.map(i => `${i.name} (${i.shape.join(',')})`).join(', ')}], ` +
            `memories: [${this.memories.map(m => `${m.input} ${m.shape} ${m.output}`).join(', ')}], ` +
            `outputs: [${this.outputs.join(', ')}] ` +
            `\n${this.layers.length} layers, ${totalUniqueWeights.toLocaleString()} weights: \n${this.layers.map(l => `${l.type} (${l})`).join('\n')}`;
    }

    public ConvertWeights(type: DataType): void {
        // for (const layer of this.layers) {
        //     if (layer.weights && layer.weights.Type !== type) {
        //         const sourceWeights = layer.weights;
        //         const targetWeights = new BarracudaArray(layer.weights.Length, type);
        //         BarracudaArray.Copy(sourceWeights, targetWeights);
        //         layer.weights = targetWeights;
        //     }
        // }
    }
}

export class ModelMetadataExtensions {
    static getTensorByName(model: Model, name: string): Tensor {
        for (const l of model.layers) {
            for (const ds of l.datasets) {
                if (ds.name === name) {
                    const tensor = new Tensor();
                    tensor.name = ds.name;
                    tensor.shape = ds.shape;
                    // tensor.cache = l.weights;
                    return tensor
                }
            }
        }
        return new Tensor();
    }

    static getShapeByName(model: Model, name: string): TensorShape | null {
        // for (const i of model.inputs) {
        //     if (i.name === name) {
        //         return new TensorShape(i.shape);
        //     }
        // }

        // let shape: TensorShape;
        // if (ModelAnalyzer.tryGetOutputTensorShape(model, name, shape)) {
        //     return shape;
        // }

        // for (const l of model.layers) {
        //     for (const ds of l.datasets) {
        //         if (ds.name === name) {
        //             return ds.shape;
        //         }
        //     }
        // }

        // for (const mem of model.memories) {
        //     if (mem.input === name || mem.output === name) {
        //         return mem.shape;
        //     }
        // }

        throw new Error("Shape " + name + " not found!");
    }

    static getDownStreamLayersCount(model: Model, name: string): number {
        return model.layers.filter(x => x.inputs.includes(name)).length;
    }
}